#ifndef PCPS_PROJECTED_OPTIMIZATION_TI_HEADER
#define PCPS_PROJECTED_OPTIMIZATION_TI_HEADER

#include <src/pcps.h>
#include <src/input.h>
#include <src/correlator.h>
#include <src/wavefunction.h>
#include <src/operator.h>
#include <src/cluster.h>
#include <src/rdm.h>
#include <src/blas_lapack.h>

namespace pcps {

  //--------------------------------------------------------------------------------
  // pcps::projected_driver_ti -- performs a projected cps optimization of the
  //                              wavefunction using translational invariance
  //--------------------------------------------------------------------------------

  template <class SCALAR>
  void projected_driver_ti(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, pcps::Wavefunction<SCALAR> & wfn);

  //-------------------------------------------------------------------------------
  // pcps::TiProjCluster -- class for a cluster of sites formed from a certain
  //                        positioning of the the projector corresponding to the
  //                        translationally invariant correlator with respect to
  //                        the cluster formed by the off-diagonal operator
  //-------------------------------------------------------------------------------

  template <class SCALAR> class TiProjCluster {

    // private data
    private:
      const int _configs_per_site;               // the number of configurations on a site
      const int _nvar;                           // the number of variables in the wavefunction
      const int _n_inner_config;                 // number of possible configurations for the inner sites
      const int _n_total_outer_occs;             // number of different total occupations the outer sites can have
      const int _jac_row_dim;                    // number of elements we will store from each jacobian row
      const int _offset;                         // position of this projectors' variables in the residual vector
      const pcps::Correlator<SCALAR> _proj_corr; // an image of the translationally invariant variable correlator
      const std::vector<int> _inner_sites;       // sites the projector shares with the operator cluster
      const std::vector<int> _outer_sites;       // projector sites that are outside the operator cluster
      std::vector<SCALAR> _rdm;                  // density matrix for different total occupations of the combined cluster
      std::vector<SCALAR> _res_data;             // accumulated data for the newton raphson residual
      std::vector<SCALAR> _jac_data;             // accumulated data for the newton raphson jacobian

      // variables are organized as follows for indexing with respect to the projection:
      //
      //    inner_cmpd * _n_total_outer_occs + outer_occ
      //
      // with the definitions:
      //
      //       inner_cmpd               the compound index for the inner sites' occupations
      //       outer_occ                the total occupation of the outer sites

    // constructor
    public:
      TiProjCluster(const pcps::Input & userinp,
                    const pcps::Wavefunction<SCALAR> & wfn,
                    const int offset,
                    const pcps::Correlator<SCALAR> & proj_corr,
                    const std::vector<int> & cluster_sites,
                    const std::vector<int> & inner_sites,
                    const std::vector<int> & outer_sites)
        : _configs_per_site(userinp.configs_per_site()),
          _nvar(wfn.nvars()),
          _n_inner_config(pcps::pow<int>(userinp.configs_per_site(), inner_sites.size())),
          _n_total_outer_occs(outer_sites.size()+1),
          _jac_row_dim(_n_inner_config * _n_total_outer_occs),
          _offset(offset),
          _proj_corr(proj_corr),
          _inner_sites(inner_sites),
          _outer_sites(outer_sites),
          _rdm(outer_sites.size() + cluster_sites.size() + 1, pcps::zero<SCALAR>()),
          _res_data(_jac_row_dim, pcps::zero<SCALAR>()),
          _jac_data(_nvar * _jac_row_dim, pcps::zero<SCALAR>())
      {

        // get the set of all sites in the cluster and projector
        std::set<int> all_sites;
        all_sites.insert(cluster_sites.begin(), cluster_sites.end());
        all_sites.insert(outer_sites.begin(), outer_sites.end());

        // create the cluster of all sites
        pcps::Cluster cluster(userinp.nsites(), userinp.configs_per_site(), all_sites);

        // get a compact RDM for the cluster of all sites
        pcps::Compact_RDM<SCALAR, pcps::OpRL<SCALAR>, pcps::Part_Spin, pcps::Ref_Uniform> rdm(userinp, wfn, pcps::OpRL<SCALAR>(), cluster);

        // extract the RDM elements
        for (int i = 0; i < _rdm.size(); i++)
          _rdm.at(i) = rdm.element(i);

      }

    // public member functions
    public:

      // function to compute a compound index from the occupations of the projected sites that lie within the operator cluster
      int inner_cmpd(const pcps::Cluster & occs) const {
        int retval = 0;
        for (int i = 0; i < _inner_sites.size(); i++)
          retval = _configs_per_site * retval + occs.ket_occ(_inner_sites[i]);
        return retval;
      }

      // function to compute the total occupation of the outer sites
      int outer_occ(const pcps::Cluster & occs) const {
        int retval = 0;
        for (int i = 0; i < _outer_sites.size(); i++)
          retval += occs.ket_occ(_outer_sites[i]);
        return retval;
      }

      // function to reset the residual and jacobian data to zero
      void reset() {
        pcps::xscal(_res_data.size(), pcps::zero<SCALAR>(), &_res_data.at(0), 1);
        pcps::xscal(_jac_data.size(), pcps::zero<SCALAR>(), &_jac_data.at(0), 1);
      }

      // function to compute the residual and jacobian contributions for a given cluster configuration
      void compute_contributions(const pcps::Input & userinp,
                                 const pcps::Cluster & bra_occs,
                                 const pcps::Cluster & ket_occs,
                                 const int total_cluster_occ,
                                 const SCALAR cluster_coeff,
                                 const std::vector<int> & bra_corr_cmp,
                                 const std::vector<int> & ket_corr_cmp,
                                 const std::vector<SCALAR> & bra_corr_val,
                                 const std::vector<SCALAR> & ket_corr_val) {

        // get the compound index for the occupation of the inner sites
        const int icmpd = this->inner_cmpd(bra_occs);

        // compute residual contributions for all total occupations of the outer sites
        pcps::xaxpy(_n_total_outer_occs, cluster_coeff, &_rdm[total_cluster_occ], 1, &_res_data[icmpd*_n_total_outer_occs], 1);
        //for (int i = 0; i < _n_total_outer_occs; i++)
        //  _res_data[ icmpd * _n_total_outer_occs + i ] += cluster_coeff * _rdm[ total_cluster_occ + i ];

        // loop over correlators in cluster
        for (int i = 0; i < ket_corr_cmp.size(); i++) {

          // compute jacobian numerator contribution
          pcps::xaxpy(_n_total_outer_occs, cluster_coeff / ket_corr_val[i],
                      &_rdm[total_cluster_occ], 1,
                      &_jac_data[ ket_corr_cmp[i] * _jac_row_dim + icmpd * _n_total_outer_occs ], 1);

          // compute jacobian denomenator contribution
          pcps::xaxpy(_n_total_outer_occs, -cluster_coeff / bra_corr_val[i],
                      &_rdm[total_cluster_occ], 1,
                      &_jac_data[ bra_corr_cmp[i] * _jac_row_dim + icmpd * _n_total_outer_occs ], 1);

        }

      }

      // function to add the residual and jacobian contributions to the total
      void add_contributions(const pcps::Input & userinp, const SCALAR coeff, std::vector<SCALAR> & res, std::vector<SCALAR> & jac) {

        // loop over the configurations of the projector
        for (pcps::Cluster pocc(userinp.nsites(), userinp.configs_per_site(), _proj_corr.site_vec()); !pocc.finished(); pocc++) {

          // get projector compound index
          const int pcmpd = _offset + _proj_corr.get_cmpd(pocc);

          // get compact residual lookup index
          const int rindex = this->inner_cmpd(pocc) * _n_total_outer_occs + this->outer_occ(pocc);

          // add residual contribution
          res[pcmpd] += coeff * _res_data[rindex];

          // add jacobian contribution
          pcps::xaxpy(_nvar, coeff, &_jac_data[rindex], _jac_row_dim, &jac[pcmpd*_nvar], 1);

        }

      }

  };

  //-------------------------------------------------------------------------------
  // pcps::TiOpCluster -- class for a cluster of correlators and sites around an
  //                      off-diagonal operator for a translationally invariant
  //                      problem
  //-------------------------------------------------------------------------------

  template <class SCALAR> class TiOpCluster {

    // private data
    private:
      std::vector<int> _offsets; // index offsets for the cluster's variable correlators
      std::vector<pcps::Correlator<SCALAR> > _var_correlators;   // TI images of the variable correlator
      std::vector<pcps::Correlator<SCALAR> > _fix_correlators;   // non-variable correlators
      pcps::Cluster _bra_occs; // bra occupation of the cluster's sites
      pcps::Cluster _ket_occs; // ket occupation of the cluster's sites
      std::vector<SCALAR> _rdm; // density matrix for different total occupations of the cluster
      std::vector<boost::shared_ptr<pcps::TiProjCluster<SCALAR> > > _proj_clusters;
      std::vector<int> _n_disjoint; // the number of disjoint projection clusters for each correlator
      int _n_var_corr;
      std::vector<int> _bra_corr_cmp; // correlator compound indices from bra occupation
      std::vector<int> _ket_corr_cmp; // correlator compound indices from ket occupation
      std::vector<SCALAR> _bra_corr_val; // correlator element values from bra occupation
      std::vector<SCALAR> _ket_corr_val; // correlator element values from ket occupation
      SCALAR _energy; // contribution to the energy
      std::vector<SCALAR> _energy_deriv; // contribution to the energy derivatives

    // constructor
    public:
      TiOpCluster(const pcps::Input & userinp,
                  const int lx,
                  const int ly,
                  const typename pcps::OpBase<SCALAR>::const_ptr & op,
                  pcps::Wavefunction<SCALAR> & wfn)
        : _energy(pcps::zero<SCALAR>()),
          _energy_deriv(wfn.nvars(), pcps::zero<SCALAR>()),
          _n_var_corr(0)
      {

        // get the number of variable correlators
        for (int c = 0; c < wfn.correlators().size(); c++)
          if (!wfn.correlators().at(c).fixed())
            _n_var_corr++;

        // for each position of the translationally invariant correlators, check if they are in the cluster
        for (int i = 0; i < lx; i++)
        for (int j = 0; j < ly; j++)
        for (int c = 0; c < wfn.correlators().size(); c++) {

          // skip fixed correlators
          if (wfn.correlators().at(c).fixed()) continue;

          // skip translations that violoate the triangular sublattice pattern
          if ( userinp.lattice_type() == pcps::keyword::Lattice_Triangular && (i+j) % 3 != 0 ) continue;

          // shift the correlator's sites
          std::vector<int> corr_sites(wfn.correlators().at(c).site_vec());
          for (std::vector<int>::iterator site = corr_sites.begin(); site != corr_sites.end(); site++) {
            const int px = *site % lx;
            const int py = *site / lx;
            *site = ( py + j ) % ly * lx + ( px + i ) % lx;
          }

          // check if the shifted correlator is touched by the operator
          bool touching = false;
          for (int k = 0; k < op->sites().size(); k++)
            if ( op->offsets().at(k) != 0 && std::count(corr_sites.begin(), corr_sites.end(), op->sites().at(k)) > 0 )
              touching = true;

          // if the correlator is touched, add it to the cluster's correlators
          if (touching) {
            _var_correlators.push_back( pcps::Correlator<SCALAR>(wfn.correlators().at(c), corr_sites) );
            _offsets.push_back(wfn.corr_offsets().at(c));
          }

        }

        // get the cluster's sites
        std::set<int> cluster_sites;
        for (int k = 0; k < op->sites().size(); k++)
          cluster_sites.insert(op->sites().at(k));
        for (typename std::vector<pcps::Correlator<SCALAR> >::const_iterator corr = _var_correlators.begin();
                                                                             corr != _var_correlators.end(); corr++)
          cluster_sites.insert(corr->site_vec().begin(), corr->site_vec().end());

        // get all fixed correlators contained by the cluster
        for (typename std::vector<pcps::Correlator<SCALAR> >::iterator corr = wfn.correlators().begin();
                                                                       corr != wfn.correlators().end(); corr++)
          if (corr->fixed()) {
            bool contained = true;
            for (std::vector<int>::const_iterator site = corr->site_vec().begin(); site != corr->site_vec().end(); site++)
              if (cluster_sites.count(*site) == 0)
                contained = false;
            if (contained)
              _fix_correlators.push_back(*corr);
          }

        // initialize the vectors holding bra and ket correlator compound indices and element values
        _bra_corr_cmp.assign(_var_correlators.size(), 0);
        _ket_corr_cmp.assign(_var_correlators.size(), 0);
        _bra_corr_val.assign(_var_correlators.size(), pcps::zero<SCALAR>());
        _ket_corr_val.assign(_var_correlators.size(), pcps::zero<SCALAR>());

        // initialize the objects used for looping over the cluster's bra and ket occupations
        _bra_occs = pcps::Cluster(userinp.nsites(), userinp.configs_per_site(), cluster_sites);
        _ket_occs = pcps::Cluster(userinp.nsites(), userinp.configs_per_site(), cluster_sites);
        for (int k = 0; k < op->sites().size(); k++)
          if        (typeid(*op) == typeid(pcps::OpNN<SCALAR>)) {
            _bra_occs.fix_occ(op->sites().at(0), 1);
            _bra_occs.fix_occ(op->sites().at(1), 1);
            _ket_occs.fix_occ(op->sites().at(0), 1);
            _ket_occs.fix_occ(op->sites().at(1), 1);
          } else if (typeid(*op) == typeid(pcps::OpN<SCALAR>)) {
            _bra_occs.fix_occ(op->sites().at(0), 1);
            _ket_occs.fix_occ(op->sites().at(0), 1);
          } else if (typeid(*op) == typeid(pcps::OpRL<SCALAR>)) {
            _bra_occs.fix_occ(op->sites().at(0), 1);
            _bra_occs.fix_occ(op->sites().at(1), 0);
            _ket_occs.fix_occ(op->sites().at(0), 0);
            _ket_occs.fix_occ(op->sites().at(1), 1);
          } else if (typeid(*op) == typeid(pcps::OpIdentity<SCALAR>)) {
          } else throw pcps::Exception( (boost::format("unknown operator type in TiOpCluster constructor")).str() );

        // compute RDM elements
        {
          _rdm.resize(cluster_sites.size()+1);
          if (cluster_sites.size() > 0) {
            pcps::Compact_RDM<SCALAR,pcps::OpRL<SCALAR>,pcps::Part_Spin,pcps::Ref_Uniform> rdm(userinp, wfn,
                                                                                               pcps::OpRL<SCALAR>(), _bra_occs);
            for (int i = 0; i < _rdm.size(); i++)
              _rdm.at(i) = rdm.element(i);
          } else
            _rdm.at(0) = pcps::unity<SCALAR>();
        }

        // initialize the different projection clusters (the first _n_var_corr clusters are for disjoint projection and operator clusters)
        _n_disjoint.resize(_n_var_corr, 0);
        _proj_clusters.resize(_n_var_corr);
        for (int i = 0; i < lx; i++)
        for (int j = 0; j < ly; j++)
        for (int c = 0, i_disjoint = 0; c < wfn.correlators().size(); c++) {

          // skip fixed correlators
          if (wfn.correlators().at(c).fixed()) continue;

          // skip translations that violoate the triangular sublattice pattern
          if ( userinp.lattice_type() == pcps::keyword::Lattice_Triangular && (i+j) % 3 != 0 ) continue;

          // shift the correlator's sites
          std::vector<int> corr_sites(wfn.correlators().at(c).site_vec());
          for (std::vector<int>::iterator site = corr_sites.begin(); site != corr_sites.end(); site++) {
            const int px = *site % lx;
            const int py = *site / lx;
            *site = ( py + j ) % ly * lx + ( px + i ) % lx;
          }

          // populate the inner and outer site vectors
          std::vector<int> inner_sites, outer_sites;
          for (std::vector<int>::const_iterator site = corr_sites.begin(); site != corr_sites.end(); site++)
            if ( cluster_sites.count(*site) > 0 )
              inner_sites.push_back(*site);
            else
              outer_sites.push_back(*site);

          // disjoint clusters
          if (outer_sites.size() == corr_sites.size()) {

            _n_disjoint.at(i_disjoint)++;
            if (_n_disjoint.at(i_disjoint) == 1)
              _proj_clusters.at(i_disjoint) = boost::shared_ptr<pcps::TiProjCluster<SCALAR> >
                                      ( new pcps::TiProjCluster<SCALAR>(userinp, wfn, wfn.corr_offsets().at(c),
                                                                        pcps::Correlator<SCALAR>(wfn.correlators().at(c), corr_sites),
                                                                        std::vector<int>(cluster_sites.begin(), cluster_sites.end()),
                                                                        inner_sites, outer_sites) );

          // touching clusters
          } else {

            _proj_clusters.push_back( boost::shared_ptr<pcps::TiProjCluster<SCALAR> >
                                       ( new pcps::TiProjCluster<SCALAR>(userinp, wfn, wfn.corr_offsets().at(c),
                                                                         pcps::Correlator<SCALAR>(wfn.correlators().at(c), corr_sites),
                                                                         std::vector<int>(cluster_sites.begin(), cluster_sites.end()),
                                                                         inner_sites, outer_sites) ) );

          }

          i_disjoint++;

        }

      }

    // public member functions
    public:

      // function to reset the energy, energy derivative, residual, and jacobian contributions, as well as the occupations
      void reset() {

        // reset energy
        _energy = pcps::zero<SCALAR>();

        // reset energy derivatives
        pcps::xscal(_energy_deriv.size(), pcps::zero<SCALAR>(), &_energy_deriv.at(0), 1);

        // reset residual projector subclusters
        for (int i = 0; i < _proj_clusters.size(); i++)
          if ( i >= _n_var_corr || _n_disjoint.at(i) > 0 )
            _proj_clusters.at(i)->reset();

        // reset occupations
        _ket_occs.reset_occs();
        _bra_occs.reset_occs();

      }

      // function to compute the energy, energy derivative, residual, and jacobian contributions for a given cluster configuration
      void compute_contributions(const pcps::Input & userinp, const pcps::Cluster & bra_occs, const pcps::Cluster & ket_occs) {

        // compute total occupation of cluster
        const int total_occ = bra_occs.total_occ();
        assert( total_occ == ket_occs.total_occ() );

        // initialize the product of correlator ratios
        SCALAR cluster_coeff = pcps::unity<SCALAR>();

        // get the bra and ket correlators' compound indices and element values
        for (int i = 0; i < _var_correlators.size(); i++) {
          _bra_corr_cmp[i] = _offsets[i] + _var_correlators[i].get_cmpd(bra_occs);
          _ket_corr_cmp[i] = _offsets[i] + _var_correlators[i].get_cmpd(ket_occs);
          _bra_corr_val[i] = _var_correlators[i][ _bra_corr_cmp[i] - _offsets[i] ];
          _ket_corr_val[i] = _var_correlators[i][ _ket_corr_cmp[i] - _offsets[i] ];
          cluster_coeff *= _ket_corr_val[i] / _bra_corr_val[i];
        }

        // get contribution from fixed correlators
        for (int i = 0; i < _fix_correlators.size(); i++)
          cluster_coeff *= _fix_correlators[i][ket_occs] / _fix_correlators[i][bra_occs];

        // add energy contribution
        _energy += cluster_coeff * _rdm[total_occ];

        // add energy derivative contributions for numerator correlators
        for (int i = 0; i < _var_correlators.size(); i++)
          _energy_deriv[_ket_corr_cmp[i]] += cluster_coeff * _rdm[total_occ] / _ket_corr_val[i];

        // add energy derivative contributions for denomenator correlators
        for (int i = 0; i < _var_correlators.size(); i++)
          _energy_deriv[_bra_corr_cmp[i]] -= cluster_coeff * _rdm[total_occ] / _bra_corr_val[i];

        // for each position of the residual projector, get its subcluster's contributions to the residual and jacobian
        for (int i = 0; i < _proj_clusters.size(); i++)
          if ( i >= _n_var_corr || _n_disjoint.at(i) > 0 )
            _proj_clusters.at(i)->compute_contributions(userinp, bra_occs, ket_occs, total_occ, cluster_coeff,
                                                        _bra_corr_cmp, _ket_corr_cmp, _bra_corr_val, _ket_corr_val);

      }

      // function to add the energy, energy derivative, residual, and jacobian contributions to the total
      void add_contributions(const pcps::Input & userinp,
                             const SCALAR coeff,
                             SCALAR & energy,
                             std::vector<SCALAR> & energy_deriv, 
                             std::vector<SCALAR> & residual, 
                             std::vector<SCALAR> & jacobian) {

        // get multiplier to deal with sublattices
        double multiplier = double(userinp.nsites());
        if ( userinp.lattice_type() == pcps::keyword::Lattice_Triangular )
          multiplier /= 3.0;

        // add energy contribution
        energy += multiplier * coeff * _energy;

        // add energy derivative contribution
        pcps::xaxpy(_energy_deriv.size(), multiplier * coeff, &_energy_deriv.at(0), 1, &energy_deriv.at(0), 1);

        // for each position of the residual projector, add the contribution to the residual and jacobian
        for (int i = 0; i < _n_var_corr; i++)
          if (_n_disjoint.at(i) > 0)
            _proj_clusters.at(i)->add_contributions(userinp, double(_n_disjoint.at(i)) * coeff, residual, jacobian);
        for (int i = _n_var_corr; i < _proj_clusters.size(); i++)
          _proj_clusters.at(i)->add_contributions(userinp, coeff, residual, jacobian);

      }

      // function to compute contributions for all cluster configurations
      void compute_all_contributions(const pcps::Input & userinp) {

        // get MPI info
        const MPI::Comm & comm = MPI::COMM_WORLD;
        const int nproc = comm.Get_size();
        const int myrank = comm.Get_rank();

        // get the total number of cluster configurations
        const int tc = _ket_occs.nconfigs();

        // get index of the starting configuration configuration for this processor
        const int extras = tc % nproc;
        const int my_min = myrank * ( tc / nproc ) + std::min(myrank, extras);
        const int my_max = ( myrank + 1 ) * ( tc / nproc ) + std::min(myrank + 1, extras);

        // convert starting index into starting configuration
        long int count = my_min;
        for (int i = _ket_occs.size()-1; i >= 0; i--) {
          if (_ket_occs.n_allowed_occs().at(i) > 1) {
            _ket_occs.set_occ(_ket_occs.sites().at(i), count % 2);
            _bra_occs.set_occ(_bra_occs.sites().at(i), count % 2);
            count = count / 2;
          }
        }

        //// print which configurations are being handled by this process
        //std::cout << boost::format("processor %2i running from %20i to %20i") % myrank % my_min % my_max << std::endl;

        // loop over all configurations of the cluster
        for (count = my_min; count < my_max; _ket_occs++, _bra_occs++, count++) {
          assert( !_bra_occs.finished() );
          assert( !_ket_occs.finished() );
          this->compute_contributions(userinp, _bra_occs, _ket_occs);
          if (myrank == 0 && userinp.verbose_print() && count % 1000000 == 0 && count > 0)
            std::cout << boost::format("processor %2i has finished %20i configurations") % myrank % count << std::endl;
        }
        if ( myrank == 0 && userinp.verbose_print() && my_max > 1000000 )
          std::cout << std::endl;

        // for clusters with no valid configurations, compute the contributions once
        if (myrank == 0 && _ket_occs.nconfigs() == 0)
          this->compute_contributions(userinp, _bra_occs, _ket_occs);

      }

  };

}

#endif
